import socket
import threading
from ollama import chat
message = [
  {
    'role': 'system',
    'content': '我是大龙虾助理，我可以这么帮到你？?',
  },
]
def generate_response(message): # binggo!!:START
    """Count=1"""
    # if count == 0:
    try:
        response = chat('qwen2.5vl', messages=message)
        return str(response['message']['content']+'——from qwen')
    except Exception as e:
        print(f"生成回复时出错: {e}")
        return "抱歉，生成回复时出错。" #Binggo!!:END
    # else:
    #     try:
    #         response = chat('qwen2.5vl', messages=message)
    #         return str(response['message']['content']+'——from qwen')
    #     except Exception as e:
    #         print(f"生成回复时出错: {e}")
    #         return "抱歉，生成回复时出错。"

def handle_client(client_socket,response): # 增加一个response参数，以确保C/S链接后，打印Ollama的回复
    """处理客户端连接"""
    print("客户端已连接")
    
    try:
        # 处理回复
        initial_message = response
        client_socket.send(initial_message.encode())
        count = 1
        while True:
            # 接收客户端消息
            message_content = client_socket.recv(4096).decode()
            if not message_content:
                break
            
            print(f"客户端: {message_content}")
            message_C2V = [
                {
                  'role' : 'user',
                  'content' : message_content, 
                },
            ]
            # 生成回复 # error
            server_response = generate_response(message_C2V)
            print(f"服务端回复: {server_response}")
            count+=1
            
            # 发送回复给客户端
            client_socket.send(server_response.encode())
            print("回复生成次数:",count)
            if input("是否继续生成回复?(y/n)") == "n":
                break
            
    except Exception as e:
        print(f"处理客户端连接时出错: {e}")
    finally:
        client_socket.close()
        print("客户端连接已关闭")

def start_server():# binggo!!:START
    """启动TCP服务器"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) #2
    server.bind(('localhost', 11435)) # 11434已经被ollama占用 #3
    server.listen(1) # 允许最多一个客户端链接 #4
    response = generate_response(message)
    # print(response) #5 #Binggo!!:END
    print("服务器已启动，等待客户端连接...") #5
    count = 0
    while True:
        client_sock, addr = server.accept()
        client_handler = threading.Thread(target=handle_client, args=(client_sock,response))
        client_handler.start()
        count += 1
        print("服务次数:",count)


if __name__ == "__main__":
    start_server() #1
